#!/bin/sh

set -eu

usage()
{
	cat << EOF

Usage: mkinitramfs [-c compression] directory output-file

Options:
  -c   Specify compression method (gzip, zstd, xz, lz4, bzip2, lzop)
       Default is gzip

This script rebuilds an initramfs from directories extracted by unmkinitramfs.
It expects to find early, early2, ... and main directories in the source directory.

EOF
}

usage_error()
{
	usage >&2
	exit 2
}

# Check if a command is available
check_command()
{
	command -v "$1" >/dev/null 2>&1 || { echo >&2 "Error: $1 is required but not installed. Aborting."; exit 1; }
}

# Create a cpio archive from a directory
create_cpio()
{
	dir="$1"
	
	if [ ! -d "$dir" ]; then
		echo "Warning: Directory $dir does not exist, skipping" >&2
		return 0
	fi
	
	(
		cd "$dir"
		find . -print0 | sort -z | cpio --null -R 0:0 -H newc -o
	)
}

# Compress a cpio archive using the specified method
compress_cpio()
{
	method="$1"
	
	case "$method" in
		gzip)
			check_command gzip
			gzip -c
			;;
		zstd)
			check_command zstd
			zstd -q -c
			;;
		xz)
			check_command xz
			xz -c
			;;
		lz4)
			check_command lz4
			lz4 -c
			;;
		bzip2)
			check_command bzip2
			bzip2 -c
			;;
		lzop)
			check_command lzop
			lzop -c
			;;
		*)
			echo "Error: Unsupported compression method: $method" >&2
			exit 1
			;;
	esac
}

# Main function to build the initramfs
build_initramfs()
{
	source_dir="$1"
	output_file="$2"
	compression="$3"
	
	# Check if source directory exists
	if [ ! -d "$source_dir" ]; then
		echo "Error: Source directory $source_dir does not exist" >&2
		exit 1
	fi
	
	# Create a temporary file for the output
	temp_file=$(mktemp "${TMPDIR:-/var/tmp}/mkinitramfs_XXXXXX")
	trap 'rm -f "$temp_file"' EXIT
	
	# Process early directories first (uncompressed cpio archives)
	# Start with early, then early2, early3, etc.
	if [ -d "$source_dir/early" ]; then
		create_cpio "$source_dir/early" >> "$temp_file"
		
		# Check for additional early directories
		i=2
		while [ -d "$source_dir/early$i" ]; do
			create_cpio "$source_dir/early$i" >> "$temp_file"
			i=$((i + 1))
		done
	fi
	
	# Process main directory (compressed cpio archive)
	if [ -d "$source_dir/main" ]; then
		create_cpio "$source_dir/main" | compress_cpio "$compression" >> "$temp_file"
	else
		# If no main directory, check if we're dealing with a simple structure
		# where all files are directly in the source directory
		if [ ! -d "$source_dir/early" ] && [ ! -d "$source_dir/early2" ]; then
			create_cpio "$source_dir" | compress_cpio "$compression" >> "$temp_file"
		fi
	fi
	
	# Move the temporary file to the output file
	mv "$temp_file" "$output_file"
	chmod 644 "$output_file"
	
	echo "Successfully created initramfs: $output_file"
}

# Parse command line options
compression="gzip"

while getopts ":c:h" opt; do
	case $opt in
		c)
			compression="$OPTARG"
			;;
		h)
			usage
			exit 0
			;;
		\?)
			echo "Invalid option: -$OPTARG" >&2
			usage_error
			;;
		:)
			echo "Option -$OPTARG requires an argument." >&2
			usage_error
			;;
	esac
done

# Shift the options so that the remaining arguments are available
shift $((OPTIND - 1))

# Check if we have the correct number of arguments
if [ $# -ne 2 ]; then
	usage_error
fi

source_dir="$1"
output_file="$2"

# Check for required commands
check_command cpio
check_command find
check_command sort

# Build the initramfs
build_initramfs "$source_dir" "$output_file" "$compression"